import os
import sys
import zmq
import subprocess
import time
import json
import cv2
import platform

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from log_info import logger

class Client():
    def __init__(self, name):
        self.client_list = {"lane": {"infer_back_dir": "infer_exe/back_infer/lane.py", "port": 5001, "size": (128, 128)},
                            "det_task": {"infer_back_dir": "infer_exe/back_infer/det_task.py", "port": 5002, "size": (416, 416)}}
        self.client_params = self.client_list[name]

        if 'Windows' in platform.platform():
            cmd_str = f'python {self.client_params["infer_back_dir"]}'  # 运行后台处理程序
            # shell=True告诉subprocess模块在运行命令时使用系统的默认shell。这使得可以像在命令行中一样执行命令，包括使用通配符和其他shell特性
            self.proc = subprocess.Popen(cmd_str)
        else:
            cmd_str = f'python3 {self.client_params["infer_back_dir"]} &'  # ubuntu和windows格式不一样
            self.proc = subprocess.Popen(cmd_str, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

        self.server = self.get_server()

    # 运行推理程序
    def init(self):
        while True:
            if self.check_ok():  # 检查后台推理程序是否ok
                break
            # 输出一个提示信息，不换行
            print('.', end='', flush=True)
            time.sleep(1)
        print(f"连接服务器{self.client_params['port']}成功")

    def __call__(self, img, timeout=10):
        self.server.setsockopt(zmq.RCVTIMEO, 1000 * timeout)
        img = cv2.resize(img, self.client_params["size"])
        img = cv2.imencode('.jpg', img)[1].tobytes()  # img转jpeg格式压缩，并转为bytes
        data = bytes('image', encoding='utf-8') + img
        self.server.send(data)
        try:
            response = self.server.recv().decode('utf-8')  # 接收服务器数据
        except self.server.timeout:
            print("接收数据超时")
            return None
        # print(response)
        response = json.loads(response)  # 把bytes转为json
        return response

    # 端口绑定,传输信息的通道
    def get_server(self):
        context = zmq.Context()
        socket = context.socket(zmq.REQ)
        socket.connect(f"tcp://127.0.0.1:{self.client_params['port']}")
        return socket

    # 检测后台程序是否成功初始化运行
    def check_ok(self):
        data = bytes('ATATA', encoding='utf-8')
        self.server.send(data)
        response = self.server.recv()  # 接收服务器数据
        response = json.loads(response)
        return response

    def close(self):
        data = bytes('Astop', encoding='utf-8')
        self.server.send(data)
        self.server.close()
        logger.info(f"{self.client_params['infer_back_dir']} close")

